#include <cuda_runtime.h>

#include <algorithm>
#include <iostream>

#include "container.cuh"

// 基本类型定义
using Shape = FixedVector<size_t>;
using Index = FixedVector<size_t>;
using SliceIndex = FixedVector<Range>;
using Axis = FixedVector<size_t>;

__device__ String ShapeToString(const Shape &shape) {
    String output;
    output.append("[");
    for (size_t i = 0; i < shape.size(); i++) {
        auto num_str = size_t_to_string(shape[i]);
        output.append(num_str.c_str());
        if (i < shape.size() - 1) output.append(", ");
    }
    output.append("]");
    return output;
}

class NdArray;
__device__ String NdArrayToString(const NdArray &x);

class NdArray {
public:
    // 简化迭代器
    using Iter = float *;
    using ConstIter = const float *;

    __device__ String ToString() {
        String output;
        auto shape_str = ShapeToString(this->m_shape);
        output.append("shape: ");
        output.append(shape_str.c_str());
        output.append("\n");
        auto data_str = NdArrayToString(*this);
        output.append(data_str.c_str());
        return output;
    }

    static __device__ NdArray from_data_and_shape(float *data, const Shape &shape) {
        NdArray arr;
        arr.set_data(data);
        arr.set_shape(shape);
        return arr;
    }
    NdArray() = default;
    NdArray(const NdArray &) = default;
    NdArray &operator=(const NdArray &) = default;

    // 基本属性
    __device__ bool empty() const {
        return size() == 0;
    }
    __device__ size_t size() const {
        size_t s = 1;
        for (size_t idx = 0; idx < m_shape.size(); idx++) s *= m_shape[idx];
        return s;
    }
    __device__ const Shape &shape() const {
        return m_shape;
    }
    __device__ size_t ndim() const {
        return m_shape.size();
    }

    __device__ float *data() {
        return m_data;
    }
    __device__ const float *data() const {
        return m_data;
    }
    // 迭代器
    __device__ Iter begin() {
        return m_data;
    }
    __device__ Iter end() {
        return m_data + size();
    }
    __device__ ConstIter begin() const {
        return m_data;
    }
    __device__ ConstIter end() const {
        return m_data + size();
    }

    // 元素访问
    __device__ float *operator()(size_t i) {
        return m_data + i;
    }

    __device__ const float *operator()(size_t i) const {
        return m_data + i;
    }

    __device__ float *operator()(const Index &index) {
        return m_data + flat_index(index);
    }

    __device__ const float *operator()(const Index &index) const {
        return m_data + flat_index(index);
    }

    // 形状操作
    __device__ NdArray reshape(const Shape &new_shape) const {
        size_t new_size = 1;
        for (size_t idx = 0; idx < new_shape.size(); idx++) new_size *= m_shape[idx];
        // assert(new_size == size() && "Reshape size mismatch");
        return NdArray::from_data_and_shape(m_data, new_shape);
    }

    __device__ void set_data(float *data) {
        m_data = data;
    }
    __device__ void set_shape(const Shape &shape) {
        m_shape = shape;
    }

    // 计算多维索引的扁平索引
    __device__ size_t flat_index(const Index &index) const {
        // assert(index.size() == m_shape.size() && "Index dimension mismatch");
        size_t idx = 0;
        size_t stride = 1;

        for (size_t i = m_shape.size(); i-- > 0;) {
            idx += index[i] * stride;
            stride *= m_shape[i];
        }

        return idx;
    }

    float *m_data;  // 外部内存指针
    Shape m_shape;  // 数组形状
};

__device__ size_t flat_index(const Index &index, const Shape &shape) {
    // assert(index.size() == m_shape.size() && "Index dimension mismatch");
    size_t idx = 0;
    size_t stride = 1;

    for (size_t i = shape.size(); i-- > 0;) {
        idx += index[i] * stride;
        stride *= shape[i];
    }

    return idx;
}

__device__ FixedVector<size_t> ComputeChildSizes(const Shape &shape) {
    const size_t n_shape = shape.size();
    if (n_shape == 0) {
        return {};
    }
    // Compute child sizes from back (the number of children for each dimension)
    FixedVector<size_t> child_sizes(n_shape, 1);
    size_t size = 1;
    for (size_t depth = n_shape - 1; 0 < depth; depth--) {
        child_sizes[depth] = size;
        size *= shape[depth];
    }
    child_sizes[0] = size;
    return child_sizes;
}

// ----------------------- Utilities for NdArray (Prsize_t) -----------------------
__device__ String ArrayLineToString(const NdArray::ConstIter &data, const size_t size) {
    String output;
    output.append("[");  // Begin of a line
    for (size_t i = 0; i < size; i++) {
        auto num_str = convert_float_to_string(data[i]);  // Output an element
        output.append(num_str.c_str());
        if (i == size - 1) {
            output.append("]");  // End of a line
        } else {
            output.append(", ");  // Splitter of an element
        }
    }
    return output;
}

__device__ String ArrayMultiDimToString(const NdArray::ConstIter &data, const Shape &shape,
                                        const FixedVector<size_t> &child_sizes, size_t depth) {
    String output;
    for (size_t i = 0; i < shape[depth]; i++) {
        // Heading
        if (i == 0) {
            output.append("[");  // begin of array
        } else {
            for (size_t d = 0; d < depth + 1; d++) {  // array indent
                output.append(" ");
            }
        }

        // Output size_ternal array
        const size_t &child_size = child_sizes[depth];
        if (depth == shape.size() - 2) {
            auto line_str = ArrayLineToString(data + child_size * i, shape[depth + 1]);
            output.append(line_str.c_str());
        } else {
            auto multi_dim_str = ArrayMultiDimToString(data + child_size * i, shape, child_sizes, depth + 1);
            output.append(multi_dim_str.c_str());
        }

        // Tailing
        if (i == shape[depth] - 1) {
            output.append("]");  // End of array
        } else {
            output.append(",\n");  // Splitter of array
        }
    }
    return output;
}

__device__ String NdArrayToString(const NdArray &x) {
    String output;
    const size_t size = static_cast<size_t>(x.size());
    const Shape &shape = x.shape();
    const FixedVector<size_t> &child_sizes = ComputeChildSizes(shape);

    if (size == 0 || shape.size() == 0) {
        // Empty
        output.append("[]");
    } else if (shape.size() == 1) {
        // 1-dim
        auto line_str = ArrayLineToString(x.data(), size);
        output.append(line_str.c_str());
    } else {
        // Multi-dim
        auto multi_dim_str = ArrayMultiDimToString(x.data(), shape, child_sizes, 0);
        output.append(multi_dim_str.c_str());
    }
    return output;
}
